import json
from typing import Callable, Dict, Literal, Optional, Union

from dbgpt._private.config import Config
from dbgpt.agent.plugin.commands.command_mange import ApiCall
from dbgpt.util.json_utils import find_json_objects

from ...memory.gpts_memory import GptsMemory
from ..agent import Agent, AgentContext
from ..base_agent import ConversableAgent

# TODO: remove global config
CFG = Config()


class DashboardAssistantAgent(ConversableAgent):
    """(In preview) Assistant agent, designed to solve a task with LLM.

    AssistantAgent is a subclass of ConversableAgent configured with a default system message.
    The default system message is designed to solve a task with LLM,
    including suggesting python code blocks and debugging.
    `human_input_mode` is default to "NEVER"
    and `code_execution_config` is default to False.
    This agent doesn't execute code by default, and expects the user to execute the code.
    """

    DEFAULT_SYSTEM_MESSAGE = """Please read the historical messages, collect the generated JSON data of the analysis sql results, and integrate it into the following JSON format to return:
      [   
        {{
            "display_type":"The chart rendering method selected for the task 1 sql",
            "sql": "Analysis sql of the step task 1",
            "thought":"thoughts summary to say to user"
        }},   
        {{
            "display_type":"The chart rendering method selected for the  task 2 sql",
            "sql": "Analysis sql of the step task 2",
            "thought":"thoughts summary to say to user"
        }}
    ]
      Make sure the response is correct json and can be parsed by Python json.loads.
    """
    DEFAULT_DESCRIBE = "Integrate analytical data generated by data scientists into a required format for building sales reports."
    NAME = "Reporter"

    def __init__(
        self,
        memory: GptsMemory,
        agent_context: AgentContext,
        describe: Optional[str] = DEFAULT_DESCRIBE,
        max_consecutive_auto_reply: Optional[int] = None,
        human_input_mode: Optional[str] = "NEVER",
        **kwargs,
    ):
        super().__init__(
            name=self.NAME,
            memory=memory,
            describe=describe,
            system_message=self.DEFAULT_SYSTEM_MESSAGE,
            max_consecutive_auto_reply=max_consecutive_auto_reply,
            human_input_mode=human_input_mode,
            agent_context=agent_context,
            **kwargs,
        )
        self.register_reply(Agent, DashboardAssistantAgent.generate_dashboard_reply)
        self.agent_context = agent_context
        self.db_connect = CFG.LOCAL_DB_MANAGE.get_connect(
            self.agent_context.resource_db.get("name", None)
        )

    async def generate_dashboard_reply(
        self,
        message: Optional[str] = None,
        sender: Optional[Agent] = None,
        reviewer: Optional[Agent] = None,
        config: Optional[Union[Dict, Literal[False]]] = None,
    ):
        """Generate a reply using code execution."""

        json_objects = find_json_objects(message)
        plan_objects = []
        fail_reason = (
            "Please recheck your answer，no usable plans generated in correct format，"
        )
        json_count = len(json_objects)
        response_success = True
        view = None
        content = None
        if json_count != 1:
            # Answer failed, turn on automatic repair
            fail_reason += f"There are currently {json_count} json contents"
            response_success = False
        else:
            try:
                chart_objs = json_objects[0]
                content = json.dumps(chart_objs, ensure_ascii=False)
                vis_client = ApiCall()
                view = vis_client.display_dashboard_vis(
                    chart_objs, self.db_connect.run_to_df
                )
            except Exception as e:
                fail_reason += f"Return json structure error and cannot be converted to a sql-rendered chart，{str(e)}"
                response_success = False

        if not response_success:
            content = fail_reason
        return True, {
            "is_exe_success": response_success,
            "content": content,
            "view": view,
        }
